# -*- coding:utf-8 -*-  
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
import datetime
import torch
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from train import train
from test import test
import config as c
def main():

    torch.autograd.set_detect_anomaly(True)
    file_dir = "runs-works"
    if not os.path.exists(file_dir):
        os.makedirs(file_dir)
    # 获取当前日期和时间
    current_time = datetime.datetime.now()
    # 格式化日期和时间
    time_str = current_time.strftime("%Y-%m-%d_%H-%M-%S")


    # name = "pre_Brightness_sampled_1000" + "_" + time_str
    # name = "pre_Contrast_sampled_1000" + "_" + time_str
    # name = "pre_Saturation_sampled_1000" + "_" + time_str
    name = "pre_Hue_sampled_1000" + "_" + time_str
    # name = "pre_Jpeg_sampled_100" + "_" + time_str

    # name = "pre_Gaussain_sampled_1000" + "_" + time_str
    # name = "pre_Salt_sampled_1000" + "_" + time_str

    train_dir = os.path.join(file_dir, name)
    os.makedirs(train_dir, exist_ok=True)

    data_dir = "/home/zzc/simulator_noisy/sampled_data"
    # sample_data_dir = data_dir + "/sample_100"
    # sample_data_dir = data_dir + "/sample_500"
    sample_data_dir = data_dir + "/sample_1000"
    # sample_data_dir = data_dir + "/sample_5000"

    train_data_dir = sample_data_dir + "/train"  # 训练数据
    test_data_dir = sample_data_dir + "/val"     # 测试数据

    model_dir = "model"
    model_dir = os.path.join(train_dir, model_dir)
    #runs-works/name/model
    os.makedirs(model_dir, exist_ok=True)
    
    log_dir = "log"
    log_dir = os.path.join(train_dir, log_dir)
    
    os.makedirs(log_dir, exist_ok=True)

    log_file = os.path.join(log_dir, "log.csv")

    save_image_dir = "image"
    save_image_dir = os.path.join(train_dir, save_image_dir)
    os.makedirs(save_image_dir, exist_ok=True)

    # Super parameter
    image_size = 128
    train_batch_size = 16
    test_batch_size = 16
    message_length = 128
    num_epochs = 5000


    # 定义图像预处理的转换操作
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomCrop((image_size, image_size), pad_if_needed=True),
            transforms.ToTensor()
        ]),
        'test': transforms.Compose([
            transforms.CenterCrop((image_size, image_size)),
            transforms.ToTensor()
        ])
    }

    # 创建ImageFolder数据集实例
    train_dataset = ImageFolder(train_data_dir, transform=data_transforms['train'])
    test_dataset = ImageFolder(test_data_dir, transform=data_transforms['test'])

    # 创建数据加载器
    train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True, num_workers=4, drop_last=True)
    test_loader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False, num_workers=4, drop_last=True)


    train(train_loader, test_loader, message_length, num_epochs, model_dir, log_file, save_image_dir)
    # test(train_loader, test_loader, message_length, num_epochs, model_dir, log_file, save_image_dir)
if __name__ == '__main__':
    main()